#############################################################################
# GFE per period
#
# - Grouped slopes, free intercepts, both varying per period
# - Groups are fixed  across periods
#
#############################################################################

rm(list=ls())
graphics.off()

library(data.table)
library(ggplot2)
library(haven)
library(xtable)
library(fixest)
library(mvtnorm)
library(Hmisc)
library("RColorBrewer")

memory.limit(size=200000)

#############################################################################
# 1. Data and generate initial values from FE
#############################################################################

# Bring data and FE results
load("Data_new/FE_period_results.RData")
rm(list=setdiff(ls(),list("dt")))
names(dt)

# Generate state-year dummies
state_list = sort(unique(dt$state))
state_year_list = c()
for (s in seq(1,length(state_list),1)) {
  # Per state, generate one dummy per year, except for the last year
  year_now = sort(unique(dt[state==state_list[s],year]))
  for (y in seq(1,length(year_now)-1,1)) {
    dt[,vnow:=as.numeric(state==state_list[s] & year==year_now[y])]
    aux=max(dt$vnow)
    if (aux==0) {
      dt[,vnow:=NULL]
    } else if (aux!=0) {
      setnames(dt,"vnow",paste0("sy_",state_list[s],"_",year_now[y]))
      state_year_list = c(state_year_list,paste0("sy_",state_list[s],"_",year_now[y]))
    }
  }
}
dt0 = copy(dt)
rm(aux,y,s,year_now)

save.image("Data_new/GFE_period/Data_for_gfe_period.RData")

#############################################################################
# 2. Run several grouped FE period
#############################################################################

# Other parameters
Ns=3000
niter=10000
tol=1e-7
sseed=20230711
np=3

# Source functions
source("Code/0. CODE Auxiliary.R")
source("Code/8a. CODE grouped FE by period.R")

# Save
save.image("Data_new/GFE_period/Data_for_gfe_period.RData")

######################################
## Version: G=2
######################################

rm(list=ls())
load("Data_new/GFE_period/Data_for_gfe_period.RData")
G=2

save_name=paste0("G",G,"_np",np)
res=tgfe_wrap(data=dt,state_year_list=state_year_list,np=np,G=G,
              niter=niter,tol=tol,Ns=Ns,sseed=sseed,save_name=save_name,se_gamma_multi=10)
save.image(paste0("Data_new/GFE_period/tgfe_",save_name,".RData"))

######################################
## Version: G=3
######################################

rm(list=ls())
load("Data_new/GFE_period/Data_for_gfe_period.RData")
G=3

save_name=paste0("G",G,"_np",np)
res=tgfe_wrap(data=dt,state_year_list=state_year_list,np=np,G=G,
              niter=niter,tol=tol,Ns=Ns,sseed=sseed,save_name=save_name,se_gamma_multi=10)
save.image(paste0("Data_new/GFE_period/tgfe_",save_name,".RData"))

######################################
## Version: G=4
######################################

rm(list=ls())
load("Data_new/GFE_period/Data_for_gfe_period.RData")
G=4

save_name=paste0("G",G,"_np",np)
res=tgfe_wrap(data=dt,state_year_list=state_year_list,np=np,G=G,
              niter=niter,tol=tol,Ns=Ns,sseed=sseed,save_name=save_name,se_gamma_multi=10)
save.image(paste0("Data_new/GFE_period/tgfe_",save_name,".RData"))


#############################################################################
# 3. Review the results one by one, put together
#############################################################################

rm(list=ls())
names_keep = c()

######################################
## Version: G=3
######################################

save_name="G3_np3"
load(paste0("Data_new/GFE_period/tgfe_",save_name,".RData"))

# Save with a new name to not overwrite
res_G3 = res$res
dc_G3 = res$dc
time_G3 = res$time
save_name = "G3"
names_keep= c(names_keep,save_name)

######################################
## Version: G=4
######################################

save_name="G4_np3"
load(paste0("Data_new/GFE_period/tgfe_",save_name,".RData"))

# Save with a new name to not overwrite
res_G4 = res$res
dc_G4 = res$dc
time_G4 = res$time
save_name="G4"
names_keep= c(names_keep,save_name)

######################################
## Version: G=2
######################################

save_name="G2_np3"
load(paste0("Data_new/GFE_period/tgfe_",save_name,".RData"))

# Save with a new name to not overwrite
res_G2 = res$res
dc_G2 = res$dc
time_G2 = res$time
save_name="G2"
names_keep= c(names_keep,save_name)


######################################
# Keep only results and save
######################################

names_keep
names_keept = c(paste0("res_",names_keep),paste0("dc_",names_keep),paste0("time_",names_keep))
names_keept = c(names_keept,"dt","dt0","state_year_list","state_list","names_keep","per")
rm(list=setdiff(ls(),names_keept))

source("Code/8a. CODE grouped FE by period.R")
save.image("Data_new/GFE_period/GFE_periods_results.RData")

# Remove intermediate files
file.remove("Data_new/GFE_period/Gfe_res1_G2_np3.RData")
file.remove("Data_new/GFE_period/Gfe_res1_G3_np3.RData")
file.remove("Data_new/GFE_period/Gfe_res1_G4_np3.RData")


